package com.lbw.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import redis.clients.jedis.JedisPooled;

@Configuration
public class EmbeddingModelConfig {

    @Value("${spring.data.redis.host}")
    private String redisHost;
    @Value("${spring.data.redis.port}")
    private int redisPort;
    @Value("${spring.ai.vectorstore.redis.prefix}")
    private String redisPrefix;
    @Value("${spring.ai.vectorstore.redis.index-name}")
    private String redisIndexName;

    @Bean
    public JedisPooled jedisPooled() {
        return new JedisPooled(redisHost,redisPort);
    }

    /**
     * 本地部署的模型 nomic-embed-text:v1.5
     * @param ollamaEmbeddingModel
     * @return
     */
    @Bean
    public RedisVectorStore redisVectorStore(JedisPooled jedisPooled, OllamaEmbeddingModel ollamaEmbeddingModel) {
        return RedisVectorStore.builder(jedisPooled,ollamaEmbeddingModel)
                .indexName(redisIndexName)
                .prefix(redisPrefix)
                .initializeSchema(true)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .build();
    }

    /**
     * 内存向量库，本地部署的模型 nomic-embed-text:v1.5
     * @param ollamaEmbeddingModel
     * @return
     */
    @Bean
    public VectorStore ollamaVectorStore(OllamaEmbeddingModel ollamaEmbeddingModel) {
        return SimpleVectorStore.builder(ollamaEmbeddingModel).build();
    }

    /**
     * openai的模型 text-embedding-v3
     * @param openAiEmbeddingModel
     * @return
     */
    @Bean
    public VectorStore openAiVectorStore(OpenAiEmbeddingModel openAiEmbeddingModel) {
        return SimpleVectorStore.builder(openAiEmbeddingModel).build();
    }


}
